package com.by.log.filter;

import com.by.log.TraceContextUtil;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

/**
 * 分布式追踪过滤器
 * 用于在RPC调用中传递追踪信息
 *
 * 该过滤器负责在RPC调用前后管理分布式追踪上下文
 */
public class DistributedTraceFilter {
    
    /**
     * 在RPC调用前设置追踪上下文
     * @param headers RPC调用的头部信息
     * @param serviceName 服务名称
     * @param methodName 方法名称
     */
    public static void beforeInvoke(Map<String, String> headers, String serviceName, String methodName) {
        // 从请求头中提取追踪信息
        TraceContextUtil.extractTraceContext(Objects.requireNonNull(headers, "RPC headers cannot be null"));
        
        TraceContextUtil.TraceData traceData = TraceContextUtil.getCurrentTraceData();
        if (traceData.getTraceId() == null) {
            // 如果没有追踪ID，则开始一个新的追踪
            TraceContextUtil.startTrace(Objects.requireNonNull(serviceName, "Service name cannot be null"), 
                                       Objects.requireNonNull(methodName, "Method name cannot be null"));
        } else {
            // 如果已有追踪ID，则开始一个新的跨度
            TraceContextUtil.startSpan(Objects.requireNonNull(serviceName, "Service name cannot be null"), 
                                      Objects.requireNonNull(methodName, "Method name cannot be null"));
        }
    }

    /**
     * 在RPC调用后清理追踪上下文
     */
    public static void afterInvoke() {
        TraceContextUtil.clear();
    }
    
    /**
     * 获取当前追踪上下文，用于传递给下游服务
     * @return 包含追踪信息的头部映射
     */
    public static Map<String, String> getCurrentContextHeaders() {
        Map<String, String> headers = new HashMap<>();
        TraceContextUtil.injectTraceContext(Objects.requireNonNull(headers, "Trace context headers cannot be null"));
        return headers;
    }
}